#pragma once
#include <folly/MPMCQueue.h>
#include <cstdlib>

class BigBlockPool {
public:
    // 构造：每块 block_size，初始化 num_blocks 块
    BigBlockPool(size_t block_size = 128*1024*1024, size_t num_blocks = 4)
        : block_size_(block_size), queue_(num_blocks) {
        for (size_t i = 0; i < num_blocks; ++i) {
            void* ptr = allocateBlock(block_size_);
            queue_.blockingWrite(ptr); // 放入队列
        }
    }

    // 禁止拷贝
    BigBlockPool(const BigBlockPool&) = delete;
    BigBlockPool& operator=(const BigBlockPool&) = delete;

    // 分配内存块（非阻塞）
    void* allocate() {
        void* ptr = nullptr;
        if (queue_.read(ptr)) {  // 非阻塞尝试读取
            return ptr;
        }
        return nullptr;  // 没有空闲块
    }

    // 分配内存块（阻塞式）
    void* allocateBlocking() {
        void* ptr = nullptr;
        queue_.blockingRead(ptr);  // 阻塞直到获取到块
        return ptr;
    }

    // 释放内存块
    void deallocate(void* ptr) {
        if (!ptr) return;
        queue_.blockingWrite(ptr);
    }

    // 析构
    ~BigBlockPool() {
        void* ptr;
        while (queue_.read(ptr)) {
            freeBlock(ptr);
        }
    }

private:
    size_t block_size_;
    folly::MPMCQueue<void*> queue_;

    // 分配单个大块
    static void* allocateBlock(size_t size) {
#if defined(_WIN32)
        void* ptr = _aligned_malloc(size, 4096);
        if (!ptr) throw std::bad_alloc();
        return ptr;
#else
        void* ptr = nullptr;
        if (posix_memalign(&ptr, 4096, size) != 0 || !ptr) {
            throw std::bad_alloc();
        }
        return ptr;
#endif
    }

    // 释放单个大块
    static void freeBlock(void* ptr) {
#if defined(_WIN32)
        _aligned_free(ptr);
#else
        free(ptr);
#endif
    }
};
